"""
子通道几何模型模块

负责处理子通道的几何参数和网格划分，包括：
1. 子通道类型定义
2. 几何参数计算
3. 网格划分
4. 连接关系建立
"""

import numpy as np
from typing import Dict, List, Tuple, Optional
from dataclasses import dataclass

@dataclass
class SubchannelType:
    """子通道类型数据类"""
    name: str                # 类型名称
    n_rods: int             # 相邻燃料棒数
    wetted_perimeter: float # 湿周
    flow_area: float        # 流通面积
    heated_perimeter: float # 加热周长
    centroid: Tuple[float, float]  # 中心点坐标

class SubchannelGeometry:
    """子通道几何模型类"""
    
    def __init__(
        self,
        n_channels: int = 36,  # 5×5棒束的36个子通道
        n_axial: int = 20,
        channel_length: float = 3.66,
        rod_diameter: float = 0.0095,  # 9.5mm
        rod_pitch: float = 0.0126      # 12.6mm
    ):
        """
        初始化子通道几何模型
        
        参数:
            n_channels (int): 子通道数量，5×5棒束为36个
            n_axial (int): 轴向节点数
            channel_length (float): 通道长度 (m)
            rod_diameter (float): 燃料棒直径 (m)
            rod_pitch (float): 燃料棒间距 (m)
        """
        self.n_channels = n_channels
        self.n_axial = n_axial
        self.channel_length = channel_length
        self.rod_diameter = rod_diameter
        self.rod_pitch = rod_pitch
        
        # 初始化子通道类型字典
        self.channel_types: Dict[str, SubchannelType] = {}
        self._initialize_channel_types()
        
        # 初始化子通道编号和类型映射
        self.channel_type_map = self._initialize_channel_type_map()
        
        # 初始化网格和连接关系
        self.dz = channel_length / n_axial  # 轴向网格尺寸
        self.channel_connections = self._initialize_channel_connections()
        self.gap_widths = self._initialize_gap_widths()
        
    def _initialize_channel_types(self):
        """初始化子通道类型"""
        # 计算基本几何参数
        P = self.rod_pitch      # 节距
        D = self.rod_diameter   # 棒径
        
        # 内部子通道（与四根燃料棒接触，每根贡献1/4圆周）
        A_inner = P**2 - np.pi * D**2 / 4           # 流通面积
        Pw_inner = np.pi * D                        # 湿周（四个1/4圆周）
        Ph_inner = np.pi * D                        # 加热周长（四个1/4圆周）
        self.channel_types['inner'] = SubchannelType(
            name='inner',
            n_rods=4,
            wetted_perimeter=Pw_inner,
            flow_area=A_inner,
            heated_perimeter=Ph_inner,
            centroid=(P/2, P/2)
        )
        
        # 边缘子通道（与两根燃料棒接触，每根贡献1/4圆周）
        A_edge = P * P / 2 - np.pi * D**2 / 8       # 流通面积
        Pw_edge = np.pi * D / 2 + P                 # 湿周（两个1/4圆周加上一个节距的壁面）
        Ph_edge = np.pi * D / 2                     # 加热周长（两个1/4圆周）
        self.channel_types['edge'] = SubchannelType(
            name='edge',
            n_rods=2,  # 修正：实际接触2根燃料棒
            wetted_perimeter=Pw_edge,
            flow_area=A_edge,
            heated_perimeter=Ph_edge,
            centroid=(P/2, 0)
        )
        
        # 角落子通道（与一根燃料棒接触，贡献1/4圆周）
        A_corner = P**2 / 4 - np.pi * D**2 / 16     # 流通面积
        Pw_corner = np.pi * D / 4 + 2 * P           # 湿周（一个1/4圆周加上两个节距的壁面）
        Ph_corner = np.pi * D / 4                   # 加热周长（一个1/4圆周）
        self.channel_types['corner'] = SubchannelType(
            name='corner',
            n_rods=1,  # 修正：实际接触1根燃料棒
            wetted_perimeter=Pw_corner,
            flow_area=A_corner,
            heated_perimeter=Ph_corner,
            centroid=(0, 0)
        )
        
    def _initialize_channel_type_map(self) -> Dict[int, str]:
        """
        初始化5×5棒束36个子通道的类型映射
        
        返回:
            Dict[int, str]: 子通道编号到类型的映射
        """
        # 创建36个子通道的类型映射
        channel_type_map = {}
        
        # 角落子通道 (4个)
        corner_channels = [1, 6, 31, 36]
        for ch in corner_channels:
            channel_type_map[ch] = 'corner'
            
        # 边缘子通道 (16个)
        edge_channels = [2, 3, 4, 5, 7, 12, 13, 18, 19, 24, 25, 30, 32, 33, 34, 35]
        for ch in edge_channels:
            channel_type_map[ch] = 'edge'
            
        # 内部子通道 (16个)
        inner_channels = [8, 9, 10, 11, 14, 15, 16, 17, 20, 21, 22, 23, 26, 27, 28, 29]
        for ch in inner_channels:
            channel_type_map[ch] = 'inner'
            
        return channel_type_map
        
    def _initialize_channel_connections(self) -> Dict[int, List[int]]:
        """
        初始化5×5棒束36个子通道的连接关系
        
        返回:
            Dict[int, List[int]]: 子通道连接关系字典
        """
        # 创建连接关系字典
        connections = {}
        
        # 定义连接关系（基于5×5棒束的实际布局）
        connections[1] = [2, 7]           # 左上角
        connections[2] = [1, 3, 8]        # 上边缘
        connections[3] = [2, 4, 9]        # 上边缘
        connections[4] = [3, 5, 10]       # 上边缘
        connections[5] = [4, 6, 11]       # 上边缘
        connections[6] = [5, 12]          # 右上角
        
        connections[7] = [1, 8, 13]       # 左边缘
        connections[8] = [2, 7, 9, 14]    # 内部
        connections[9] = [3, 8, 10, 15]   # 内部
        connections[10] = [4, 9, 11, 16]  # 内部
        connections[11] = [5, 10, 12, 17] # 内部
        connections[12] = [6, 11, 18]     # 右边缘
        
        # ... 继续定义其他子通道的连接关系
        # 中间行的连接关系
        connections[13] = [7, 14, 19]
        connections[14] = [8, 13, 15, 20]
        connections[15] = [9, 14, 16, 21]
        connections[16] = [10, 15, 17, 22]
        connections[17] = [11, 16, 18, 23]
        connections[18] = [12, 17, 24]
        
        # ... 继续定义剩余子通道的连接关系
        
        # 最后一行的连接关系
        connections[31] = [32, 25]        # 左下角
        connections[32] = [31, 33, 26]    # 下边缘
        connections[33] = [32, 34, 27]    # 下边缘
        connections[34] = [33, 35, 28]    # 下边缘
        connections[35] = [34, 36, 29]    # 下边缘
        connections[36] = [35, 30]        # 右下角
        
        return connections
        
    def _initialize_gap_widths(self) -> Dict[Tuple[int, int], float]:
        """
        初始化子通道间隙宽度
        
        返回:
            Dict[Tuple[int, int], float]: 子通道间隙宽度字典
        """
        gap_widths = {}
        
        # 计算标准间隙宽度（燃料棒间距减去燃料棒直径）
        standard_gap = self.rod_pitch - self.rod_diameter
        
        # 为每个连接分配间隙宽度
        for channel_id, connected_channels in self.channel_connections.items():
            for connected_id in connected_channels:
                if (channel_id, connected_id) not in gap_widths and \
                   (connected_id, channel_id) not in gap_widths:
                    gap_widths[(channel_id, connected_id)] = standard_gap
                    gap_widths[(connected_id, channel_id)] = standard_gap
                    
        return gap_widths
        
    def get_channel_type(self, channel_id: int) -> str:
        """
        获取指定子通道的类型
        
        参数:
            channel_id (int): 子通道编号
            
        返回:
            str: 子通道类型名称
        """
        return self.channel_type_map[channel_id]
        
    def get_channel_params(self, channel_id: int) -> Dict[str, float]:
        """
        获取指定子通道的几何参数
        
        参数:
            channel_id (int): 子通道编号
            
        返回:
            Dict[str, float]: 子通道几何参数字典
        """
        channel_type = self.get_channel_type(channel_id)
        type_data = self.channel_types[channel_type]
        
        return {
            'flow_area': type_data.flow_area,
            'wetted_perimeter': type_data.wetted_perimeter,
            'heated_perimeter': type_data.heated_perimeter,
            'hydraulic_diameter': 4.0 * type_data.flow_area / type_data.wetted_perimeter
        }
        
    def get_connected_channels(self, channel_id: int) -> List[int]:
        """
        获取与指定子通道相连的子通道列表
        
        参数:
            channel_id (int): 子通道编号
            
        返回:
            List[int]: 相连子通道的编号列表
        """
        return self.channel_connections[channel_id]
        
    def get_gap_width(self, channel_1: int, channel_2: int) -> Optional[float]:
        """
        获取两个子通道之间的间隙宽度
        
        参数:
            channel_1 (int): 第一个子通道编号
            channel_2 (int): 第二个子通道编号
            
        返回:
            Optional[float]: 间隙宽度 (m)，如果两个子通道不相连则返回None
        """
        if (channel_1, channel_2) in self.gap_widths:
            return self.gap_widths[(channel_1, channel_2)]
        elif (channel_2, channel_1) in self.gap_widths:
            return self.gap_widths[(channel_2, channel_1)]
        return None
        
    def get_axial_position(self, node_id: int) -> float:
        """
        获取轴向节点的位置
        
        参数:
            node_id (int): 节点编号
            
        返回:
            float: 节点位置 (m)
        """
        return node_id * self.dz
        
    def get_channel_centroid(self, channel_id: int) -> Tuple[float, float]:
        """
        获取子通道的中心点坐标
        
        参数:
            channel_id (int): 子通道编号
            
        返回:
            Tuple[float, float]: 中心点坐标 (x, y)
        """
        channel_type = self.get_channel_type(channel_id)
        return self.channel_types[channel_type].centroid